package cn.qiao.mdc.thread;

import cn.qiao.mdc.consts.Constants;
import org.slf4j.MDC;

import java.util.Map;

public class RunnableWrapper implements Runnable {

    private final Runnable runnable;
    private final Map<String, String> contextMap;

    public RunnableWrapper(Runnable runnable) {
        this.runnable = runnable;
        this.contextMap = MDC.getCopyOfContextMap();
    }

    @Override
    public void run() {
        try {
            if (contextMap != null) {
                MDC.setContextMap(contextMap);
            } else {
                MDC.put(Constants.TRACE_ID, Constants.getTraceId());
            }
            if (runnable != null) {
                runnable.run();
            }
        } finally {
            MDC.clear();
        }
    }
}
